from py3gpp import *
import numpy as np
import scipy.signal as signal
from gscn import GSCN2CentFreq,gscn_pssIndices
import sigmf
from pss import PSS

class RawData:
    def __init__(self,fn) -> None:
        self.fn = fn
        self.handle = sigmf.sigmffile.fromfile(fn)

        self.capture = self.handle.get_capture_info(0)
        self.f_center = self.capture.get(sigmf.SigMFFile.FREQUENCY_KEY, 0)
        self.fs = self.handle.get_global_field(sigmf.SigMFFile.SAMPLE_RATE_KEY)

        self.f_start = self.f_center-self.fs/2
        self.f_end = self.f_center+self.fs/2

    def read(self,start=5e-3,end=30e-3):
        if end>0:
            self.waveform100 = self.handle.read_samples()[int(start * self.fs):int(end * self.fs)]
        else:
            self.waveform100 = self.handle.read_samples()[int(start * self.fs):]
        self.waveform100 /= max(self.waveform100.real.max(), self.waveform100.imag.max())
            
    def decimate(self,f):
        df = -(f-self.f_center)/self.fs
        self.waveform = signal.decimate(self.waveform100*np.exp(1j*2*np.pi*df*np.arange(len(self.waveform100))),self.dec_factor,ftype='fir')

class LTE(RawData):
    def __init__(self,fn) -> None:
        RawData.__init__(self,fn)
        self.mu = 0
        self.dec_factor = int(self.fs/(128*15e3 * 2**self.mu))

    def sync(self):
        peak_value = np.zeros(3)
        peak_index = np.zeros(3, 'int')
        for current_NID2 in np.arange(3, dtype='int'):
            pss = PSS(current_NID2)
            temp = signal.correlate(self.waveform, pss.td, 'valid')  # correlate over 25 ms
            index = np.argmax(np.abs(temp))
            peak = np.abs(temp[index])
            peak_index[current_NID2] = index
            peak_value[current_NID2] = peak

        self.detected_NID2 = np.argmax(peak_value)
        self.timingOffset = peak_index[self.detected_NID2]
        self.pss_peak =  peak_value[self.detected_NID2]
    
    def search(self,start=None,end=None):
        if start==None:
            start = self.f_center - 50e6
        if end==None:
            end = self.f_center + 50e6
        sg = int(start/100e3)
        se = int(end/100e3)
        n = se - sg
        self.offset = np.zeros((n,),dtype=np.int64)
        self.nid2 = np.zeros((n,),dtype=np.int32)
        self.peak = np.zeros((n,))

        for k in range(n):
            f = start + k*100e3
            self.decimate(f)
            self.waveform /= max(self.waveform.real.max(), self.waveform.imag.max())
            self.sync()
            self.offset[k] = self.timingOffset
            self.peak[k] = self.pss_peak
            self.nid2[k] = self.detected_NID2
        #     if k%10==0:
        #         print("*",end='')
        # print()
        k = np.argmax(self.peak)
        peak = self.peak[k]
        self.pss_all = []

        for k in range(n):
            if self.peak[k]>0.5*peak:
                self.pss_all.append((start+k*100e3,self.nid2[k],self.offset[k],self.peak[k]))
        return self.pss_all

class NR(RawData):
    def __init__(self,fn) -> None:
        RawData.__init__(self,fn)

        self.mu = 1
        self.reset()
        
    def reset(self):
        self.carrier = nrCarrierConfig(NSizeGrid = 20, SubcarrierSpacing = 15 * 2**self.mu)
        self.info = nrOFDMInfo(self.carrier)
        self.Nfft = self.info['Nfft']

        slotGrid = nrResourceGrid(self.carrier)
        self.slotGrid = slotGrid[:, 0]
        self.dec_factor = int(self.fs/(self.Nfft*15e3 * 2**self.mu))
        
        self.detected_NID = -1
        
        self.max_frames = 4
    
    def set_search_end(self,up):
        self.f_end = up
        return self
    
    def set_search_start(self,down):
        self.f_start = down
        return self
    
    def decimate(self,gscn):
        f_ssb = self.GSCN2CentFreq(gscn)
        RawData.decimate(self,f_ssb)

    def sync(self):
        peak_value = np.zeros(3)
        peak_index = np.zeros(3, 'int')
        PSS_LEN = 128
        NRE_PER_PRB = 12

        start = 56
        pssIndices = np.arange(start, start + PSS_LEN  - 1)
        SampleRate = self.fs/8

        for current_NID2 in np.arange(3, dtype='int'):
            
            self.slotGrid[pssIndices] = nrPSS(current_NID2)
            [refWaveform, info] = nrOFDMModulate(self.carrier, self.slotGrid, SampleRate = SampleRate)

            refWaveform = refWaveform[info['CyclicPrefixLengths'][0]:]; # remove CP

            temp = signal.correlate(self.waveform, refWaveform, 'valid')  # correlate over 25 ms
            index = np.argmax(np.abs(temp))
            peak = np.abs(temp[index])
            peak_index[current_NID2] = index
            peak_value[current_NID2] = peak

        self.detected_NID2 = np.argmax(peak_value)
        self.timingOffset = peak_index[self.detected_NID2]
        self.peak =  peak_value[self.detected_NID2]
    
    def GSCN2CentFreq(self,gscn):
        self.gscn = gscn
        if gscn<7499:
            n = gscn % 3
            if n==0:
                N = gscn // 3
                M = 3
            elif n==1:
                N = gscn // 3
                M = 5
            else:
                N = gscn // 3 + 1
                M = 1
            return N*1.2e6+M*50e3
        elif gscn<22256:
            return (gscn-7499)*1.44e6+3e9
        else:
            return (gscn-22256)*17.28e6+24250.08e6
        
    @staticmethod
    def nearGSCN(fc):
        if fc<3e9:
            return int(round(fc/1.2e6*3))
        elif fc<24.250e9:
            return 7499+int(round((fc-3e9)/1.44e6))
        else:
            return 22256+int(round((fc-24250.08e6)/17.28))
        
    def report(self):
        r = f'{self.fn}:GSCN={self.gscn}  detect PSS {self.detected_NID2}@{self.timingOffset} with {self.peak}'
        if self.detected_NID!=-1:
            r += f' SSS {self.detected_NID}'
        return r
    
    def search(self,start=None,end=None):
        if start==None:
            start = self.f_start
        if end==None:
            end = self.f_end
        sg = NR.nearGSCN(start)
        se = NR.nearGSCN(end)
        n = se - sg
        offset = np.zeros((n,),dtype=np.int64)
        nid2 = np.zeros((n,),dtype=np.int32)
        peak = np.zeros((n,))
        for k in range(n):
            gscn = k+sg
            self.decimate(gscn)
            self.sync()
            offset[k] = self.timingOffset
            peak[k] = self.peak
            nid2[k] = self.detected_NID2
        k = np.argmax(peak)
        self.gscn = k+sg
        self.detected_NID2 = nid2[k]
        self.peak = peak[k]
        self.timingOffset = offset[k]
        return k+sg,nid2[k],peak[k],offset[k]
    def find_psses(self,waveform,alpha):
        PSS_LEN = 128
        NRE_PER_PRB = 12
        start = 56
        pssIndices = np.arange(start, start + PSS_LEN  - 1)
        SampleRate = self.fs/8
        all_pss = []
        for nid2 in range(3):
            self.slotGrid[pssIndices] = nrPSS(nid2)
            [refWaveform, info] = nrOFDMModulate(self.carrier, self.slotGrid, SampleRate = SampleRate)
            refWaveform = refWaveform[info['CyclicPrefixLengths'][0]:]; # remove CP
            temp = signal.correlate(waveform, refWaveform, 'valid')  # correlate over 25 ms       
            index = 10
            while index < len(temp):
                if np.abs(temp[index])>self.peak*alpha:
                    l_ = np.argmax(np.abs(temp[index-5:index+5]))
                    all_pss.append((nid2, index-5+l_,temp[index-5+l_]))
                    index += 10
                else:
                    index += 1
        return all_pss
    
    def find_all_pss(self,alpha):
        self.decimate(self.gscn)
        self.all_pss = self.find_psses(self.waveform,alpha)
        return self.all_pss
    
    def sssCheck(self,time_off=None):
        if time_off==None:
            time_off = self.timingOffset
        self.decimate(self.gscn)
        tos = time_off - (127-118)*4
        rxGrid = nrOFDMDemodulate(carrier = self.carrier, waveform = self.waveform[tos:], initialNSlot = 2, SampleRate=self.fs/self.dec_factor, CyclicPrefixFraction=0.5)
        rxGrid /= np.max((rxGrid.real.max(), rxGrid.imag.max()))
        sssEst = np.zeros(336*3)
        start = 56
        PSS_LEN = 128
        sssIndices = np.arange(start, start + PSS_LEN  - 1)
        sssRx = nrExtractResources(sssIndices, rxGrid[:,2:])
        for ncellid in range(336*3):

            sssRef = nrSSS(ncellid)
            sssEst[ncellid] = np.abs(np.vdot(sssRx, sssRef))

        self.detected_NID = np.argmax(sssEst)
        return self.detected_NID
    
    def detect_NID(self,waveform,nid2,offset):
        scsSSB = 15 * 2**(self.mu)
        nSlot = 0
        rxSampleRate = self.fs/8
        nrbSSB = 20

        pss_pos = 592
        rxGrid = nrOFDMDemodulate(waveform = waveform[offset-pss_pos:][:np.min((len(waveform), 2048*8))], nrb = nrbSSB, scs = scsSSB, initialNSlot = nSlot, SampleRate=rxSampleRate, CyclicPrefixFraction=0.5)
        rxGrid = rxGrid[:,1:5]
        rxGrid /= np.max((rxGrid.real.max(), rxGrid.imag.max()))
        sssIndices = nrSSSIndices()
        sssRx = nrExtractResources(sssIndices, rxGrid)
        sssRx /= np.max((sssRx.real.max(), sssRx.imag.max())) # scale sssRx symbols individually
        
        sssEst = np.zeros(336)
        for NID1 in range(335):
            ncellid = (3*NID1) + nid2
            sssRef = nrSSS(ncellid)
            sssEst[NID1] = np.abs(np.vdot(sssRx, sssRef))
        detected_NID1 = np.argmax(sssEst)
        peak = sssEst[detected_NID1]
        detected_NID = detected_NID1*3 + nid2
        return detected_NID,peak
        
    def checkSSB(self,nid2,offset):
        return self.checkSSBwithWaveform(self.waveform,nid2,offset)
    
    def checkSSBwithWaveform(self,waveform,nid2,offset):
        scsSSB = 15 * 2**(self.mu)
        nSlot = 0
        rxSampleRate = self.fs/8
        nrbSSB = 20

        pss_pos = 592
        rxGrid = nrOFDMDemodulate(waveform = waveform[offset-pss_pos:][:np.min((len(waveform), 2048*8))], nrb = nrbSSB, scs = scsSSB, initialNSlot = nSlot, SampleRate=rxSampleRate, CyclicPrefixFraction=0.5)
        rxGrid = rxGrid[:,1:5]
        rxGrid /= np.max((rxGrid.real.max(), rxGrid.imag.max()))
        sssIndices = nrSSSIndices()
        sssRx = nrExtractResources(sssIndices, rxGrid)
        sssRx /= np.max((sssRx.real.max(), sssRx.imag.max())) # scale sssRx symbols individually
        
        sssEst = np.zeros(336)
        for NID1 in range(335):
            ncellid = (3*NID1) + nid2
            sssRef = nrSSS(ncellid)
            sssEst[NID1] = np.abs(np.vdot(sssRx, sssRef))
        detected_NID1 = np.argmax(sssEst)
        detected_NID = detected_NID1*3 + nid2
        
        dmrsIndices = nrPBCHDMRSIndices(detected_NID, style='matlab')
        xcorrPBCHDMRS = np.empty(7)
        for ibar_SSB in range(7):
            PBCHDMRS = nrPBCHDMRS(detected_NID, ibar_SSB)
            xcorrPBCHDMRS[ibar_SSB] = np.abs(np.vdot(nrExtractResources(dmrsIndices, rxGrid), PBCHDMRS))
        ibar_SSB = np.argmax(np.abs(xcorrPBCHDMRS))
        # print(f"detect {detected_NID} ibar_ssb {ibar_SSB}")
        
        refGrid = np.zeros((nrbSSB*12, 4), 'complex')
        dmrsIndices = nrPBCHDMRSIndices(detected_NID)
        temp = nrSetResources(dmrsIndices, refGrid, nrPBCHDMRS(detected_NID, ibar_SSB))
        H, nVar = nrChannelEstimate(rxGrid = rxGrid, refGrid = refGrid)
        pbchIndices = nrPBCHIndices(detected_NID)
        pbch_eqed, csi = nrEqualizeMMSE(nrExtractResources(pbchIndices, rxGrid), nrExtractResources(pbchIndices, H), nVar)
        pbchBits = nrSymbolDemodulate(pbch_eqed, 'QPSK', nVar, 'soft')
        E = 864
        v = ibar_SSB
        scrambling_seq = nrPBCHPRBS(detected_NID, v, E)
        scrambling_seq_bpsk = (-1)*scrambling_seq*2 + 1
        pbchBits_descrambled = pbchBits * scrambling_seq_bpsk
        pbchBits_csi = pbchBits_descrambled * np.repeat(csi, 2)
        decoded3, crc_result, payload, lsbotfsfn, hrf, msbidxoffset = nrBCHDecode(pbchBits_csi,8,0,detected_NID)
         
        return crc_result, detected_NID, lsbotfsfn, ibar_SSB
        
    def checkSSB_all(self):
        self.all_ssb = []
        print(f"checking SSB, total: {len(self.all_pss)}")
        for nid2,offset,_ in self.all_pss:
            crc_result, detected_NID, lsbotfsfn, _, _ = self.checkSSB(nid2,offset)
            if crc_result==0:
                self.all_ssb.append(SSBInfo().cellid(detected_NID).offset(offset-592))
        return self.all_ssb
    
    def step0(self):
        self.read(end=30e-3)
        gscn,nid2,peak,offset = self.search()
        print(gscn,nid2,peak,offset)
        self.read(end=-1.0)
        self.decimate(gscn)
        self.frames = int(len(self.waveform)/(self.fs/8*20e-3))-1
        crc_result, detected_NID, _, ibar_SSB = self.checkSSB(self.detected_NID2,self.timingOffset) 
        if crc_result==0:
            self.gscn = gscn
            self.detected_NID = detected_NID
            self.ibar_SSB = ibar_SSB
            self.synced = True
        else:
            self.synced = False
        return crc_result==0
    
    def report_nid(self):
        if self.synced:
            return f"gscn={self.gscn}, nid={self.detected_NID}, ibar={self.ibar_SSB}"
        else:
            return "no synced"
    
    def search_t(self, wv, ref):
        l = len(wv)
        res = 256
        step = 1/res
        fwv = np.fft.fft(wv)
        fref = np.fft.fft(ref)
        est = np.zeros((4*res,))
        for i in range(4*res):
            est[i] = np.abs(np.dot(fwv*np.exp(1j*2*np.pi*np.arange(l)*(i-2*res)*step/l),fref.conj()))
        i = np.argmax(est)
        return (i-2*res)*step, est
    
    def calibration_t(self, wv, t):
        l = len(wv)
        fwv = np.fft.fft(wv)
        ret = np.fft.ifft(fwv*np.exp(1j*2*np.pi*np.arange(l)*t/l))
        return ret
    
    def est(self):
        SampleRate = self.fs/8
        ss = int(round(20e-3*SampleRate))
        slotGrid = nrResourceGrid(self.carrier)
        slotGrid = slotGrid[:,:4]
        nrSetResources(nrPSSIndices(),slotGrid,nrPSS(self.detected_NID%3))
        nrSetResources(nrSSSIndices(),slotGrid,nrSSS(self.detected_NID))
        dmrsIndices = nrPBCHDMRSIndices(self.detected_NID)
        nrSetResources(dmrsIndices, slotGrid, nrPBCHDMRS(self.detected_NID, self.ibar_SSB))
        [refWaveform, info] = nrOFDMModulate(self.carrier, slotGrid, SampleRate = SampleRate)
        refWaveform = refWaveform[info['CyclicPrefixLengths'][0]:]; # remove CP
        ref_len = len(refWaveform)
        self.phase = np.zeros((self.frames,),dtype=np.complex128)
        self.dt = np.zeros((self.frames,))
        for i in range(self.frames):
            wv = self.waveform[self.timingOffset+i*ss:self.timingOffset+i*ss+ref_len]
            self.dt[i],_ = self.search_t(wv,refWaveform)
            c_wv = self.calibration_t(wv,self.dt[i])
            self.phase[i] = np.dot(c_wv,refWaveform.conj())

    def step1(self,alpha=1.0):
        SampleRate = self.fs/8
        ss = int(round(20e-3*SampleRate))
        self.est()
        ssb_len = 512*8
        self.ssbWaveform = np.zeros((2*ssb_len,),dtype=np.complex64)
        for i in range(min(self.frames,self.max_frames)):
            wv = self.waveform[self.timingOffset+i*ss-ssb_len:self.timingOffset+i*ss+ssb_len]
            c_wv = self.calibration_t(wv,self.dt[i])
            self.ssbWaveform += c_wv * self.phase[i].conj()
        self.ssbWaveform /= max(self.ssbWaveform.real.max(), self.ssbWaveform.imag.max()) 
        self.step1_pss = self.find_psses(self.ssbWaveform,alpha) 
        self.step1_pss.sort(key=lambda x:np.abs(x[2]))
        return self.step1_pss
    
    def step2(self):
        ssb_len = 512*8
        self.ssbs = []
        self.step2_sss = []
        for nid2, off, _ in self.step1_pss[:10]:
            crc_result, detected_NID, _, ibar_SSB = self.checkSSB(nid2,self.timingOffset+off-ssb_len)
            detected_NID, peak = self.detect_NID(self.ssbWaveform,nid2,off)
            self.step2_sss.append((detected_NID,np.abs(peak)))
        
    def step3(self):
        ssb_len = 512*8
        self.ssbs = []
        for nid2, off, _ in self.step1_pss[:10]:
            crc_result, detected_NID, _, ibar_SSB = self.checkSSB(nid2,self.timingOffset+off-ssb_len)
            detected_NID, peak = self.detect_NID(self.ssbWaveform,nid2,off)
            if crc_result == 0:
                detected_NID, peak = self.detect_NID(self.ssbWaveform,nid2,off)
                self.ssbs.append(SSBInfo().cellid(detected_NID).offset(self.timingOffset+off-ssb_len).set_peak(peak).set_ibar(ibar_SSB))
        return self.ssbs
          
                    
class SSBInfo:
    def __init__(self):
        pass
    def cellid(self, nid):
        self.nid = nid
        return self
    def offset(self,off):
        self.off = off
        return self
    def set_ibar(self,i):
        self.ibar = i
        return self
    def set_peak(self,peak):
        self.peak = peak
        return self
    def __str__(self):
        return f"NID {self.nid} @{self.off} peak={self.peak} ibar={self.ibar}"